//
// Created by Alex on 17/3/8.
//

#include "uthread.h"
#include <stdio.h>

static thread_local uthread_t *g_callback[128];
static thread_local int g_callbackCnt = 0;

stack_mem_t::stack_mem_t()
        : m_occupy(nullptr), m_stackMem(nullptr), m_stackSize(0) {}

stack_mem_t::stack_mem_t(int32_t size)
        : m_occupy(nullptr), m_stackMem(new char[size]), m_stackSize(size) {}

void uthread_t::saveStack() {
    size_t len =
            (size_t) ((intptr_t) m_context.fc_stack.base - (intptr_t) m_stackTop);
    if (m_saveStack != nullptr && len > m_saveSize) {
        delete[] m_saveStack;
        m_saveStack = nullptr;
    }
    if (m_saveStack == nullptr) {
        m_saveStack = new char[len];
        m_saveSize = len;
    }
    memcpy(m_saveStack, m_stackTop, m_saveSize);
}

uthread_t::~uthread_t() {
    if (m_saveStack != nullptr) {
        delete[] m_saveStack;
        m_saveStack = nullptr;
        m_stackTop = nullptr;
    }
    if (!m_isShareStack && m_stack != nullptr) {
        if (m_stack->m_stackMem != nullptr) {
            delete m_stack->m_stackMem;
        }
        delete m_stack;
        m_stack = nullptr;
    }

    if (m_spec != nullptr) {
        delete m_spec;
        m_spec = nullptr;
    }
}

void uthread_t::resumeStack() { memcpy(m_stackTop, m_saveStack, m_saveSize); }

inline static int32_t getStackSize(int32_t stackSize);

share_stack_t::share_stack_t(int32_t count, int32_t stackSize)
        : m_allocId(0), m_count(count > 0 ? count : 1) {

    m_stackSize = getStackSize(stackSize);
    m_stackInfo = new stack_mem_t[count];
    m_mem = new char[m_stackSize * m_count];

    for (int i = 0; i < m_count; ++i) {
        m_stackInfo[i].m_stackMem = m_mem + i * m_stackSize;
        m_stackInfo[i].m_stackSize = m_stackSize;
    }
}

share_stack_t::~share_stack_t() {
    if (m_stackInfo != nullptr) {
        delete[] m_stackInfo;
        m_stackInfo = nullptr;
    }
    if (m_mem != nullptr) {
        delete[] m_mem;
        m_mem = nullptr;
    }
}

stack_mem_t *share_stack_t::allocStack() {
    stack_mem_t *ret = &m_stackInfo[m_allocId];
    m_allocId = (m_allocId + 1) % m_count;
    return ret;
}

static void uthread_run_proc(intptr_t arg) {
    uthread_t *co = (uthread_t *) arg;
    if (co != nullptr && co->m_proc != nullptr) {
        co->m_proc(co->m_arg);
        co->m_exit = true;
        co->yield();
    }
}

inline static int32_t getStackSize(int32_t stackSize) {
    if (stackSize <= 0) {
        stackSize = kDefaultStackSize;
    } else if (stackSize > kMaxStackSize) {
        stackSize = kMaxStackSize;
    }
    if (stackSize & 0xFFF) {
        stackSize &= 0xFFF;
        stackSize += 0xFFF;
    }
    return stackSize;
}

uthread_t::uthread_t(uthread_attr_t *attr, uthread_proc proc, void *arg)
        : m_proc(proc), m_arg(arg), m_context(), m_exit(false), m_active(false),
          m_isMain(false), m_enableHookSys(false), m_spec(nullptr), m_specCount(0) {
    uthread_attr_t at;
    if (attr != nullptr) {
        memcpy(&at, attr, sizeof(uthread_attr_t));
    }
    at.stack_size = getStackSize(at.stack_size);

    if (at.share_stack != nullptr) {
        m_stack = at.share_stack->allocStack();
    } else {
        m_stack = new stack_mem_t(at.stack_size);
    }
    m_isShareStack = at.share_stack != nullptr;
    m_context.fc_stack.base = m_stack->m_stackMem + m_stack->m_stackSize;
    m_context.fc_stack.limit = m_stack->m_stackMem;
}

void uthread_t::resume() {
    if (m_exit || g_callbackCnt > 0 && g_callback[g_callbackCnt - 1] == this) {
        return;
    }
    uthread_t *occupy = m_stack->m_occupy;
    bool isStackOccupy = (occupy != nullptr && occupy != this);

    if (m_isShareStack && isStackOccupy) //别的协程占用了stack, 先保存它的stack
    {
        occupy->saveStack();
    }

    if (!m_active) {
        make_fcontext(&m_context, uthread_run_proc);
        m_active = true;
    } else {
        if (m_isShareStack &&
            isStackOccupy) { //已经运行, 但占用stack的协程不是自身,需要恢复
            resumeStack();
        }
    }

    m_stack->m_occupy = this;

    g_callback[g_callbackCnt] = this;
    g_callbackCnt += 1;
    jump_fcontext(&(m_scheduleContext), &(m_context), (intptr_t) this);
}

void uthread_t::yield() {
    if (g_callbackCnt <= 0) {
        return;
    }
    uthread_t *running = g_callback[g_callbackCnt - 1];
    char dummy;
    running->m_stackTop = &dummy;

    g_callbackCnt -= 1;
    jump_fcontext(&(running->m_context), &(running->m_scheduleContext), 0);
}

void cr_yield() {
    if (g_callbackCnt > 0) {
        g_callback[g_callbackCnt - 1]->yield();
    }
}

void cr_create(uthread_t *&ut, uthread_attr_t *attr, void (*fun)(void *),
               void *arg) {
    ut = new uthread_t(attr, fun, arg);
}

void cr_resume(uthread_t *ut) { ut->resume(); }

void cr_release(uthread_t *co) {
    if (co->m_exit) {
        delete co;
    }
}

uthread_t *cr_get_curr_uthread() {
    return g_callbackCnt > 0 ? g_callback[g_callbackCnt - 1] : nullptr;
}

share_stack_t *cr_create_share_stack(int32_t count, int32_t stack_size) {
    return new share_stack_t(count, stack_size);
}

void cr_release_share_stack(share_stack_t *share) {
    if (share != nullptr) {
        delete share;
    }
}

inline void *uthread_t::getSpec(int key) {
    int64_t val = key < m_specCount && key > 0 ? m_spec[key] : 0;
    return (void *) val;
}

inline void *uthread_t::setSpec(int key, void *val) {
    if (m_spec == nullptr || m_specCount <= 0) {
        m_spec = new uint64_t[key + 1];
        m_specCount = key + 1;
    } else if (m_specCount < key) {
        uint64_t *spec = new uint64_t[key + 1];
        memcpy(spec, m_spec, m_specCount);

        delete[] m_spec;
        m_spec = spec;
        m_specCount = key + 1;
    }

    m_spec[key] = (uint64_t) val;
}

void *cr_get_sepc(uthread_t *ut, int key) {
    return ut->getSpec(key);
}

void cr_set_spec(uthread_t *ut, int key, void *value) {
    ut->setSpec(key, value);
}